from numpy import cos, sin, pi, absolute, arange, maximum, log10
from scipy.signal import lfilter, firwin, freqz
from pylab import figure, clf, plot, xlabel, ylabel, ylim, title, grid, show, subplot, subplots_adjust

#------------------------------------------------
# Create a signal x for demonstration.
#------------------------------------------------
sample_rate = 100.0
nsamples = 400
t = arange(nsamples) / sample_rate
x = cos(2*pi*0.5*t) + \
    0.3*sin(2*pi*15.3*t) +\
    0.2*sin(2*pi*18.7*t + 0.1) + \
    0.1*sin(2*pi*25.5*t + 0.8)

#------------------------------------------------
# Create a FIR filter and apply it to x.
#------------------------------------------------
# The Nyquist rate of the signal.
nyq_rate = sample_rate / 2.0

# The cutoff frequency of the filter.
cutoff_hz = 10.0

# Fix the order for the FIR filter.
N = 74

# Compute the Taps of the filer.
taps = firwin(N, cutoff_hz/nyq_rate, window = "hamming")
 
# Use lfilter to filter x with the FIR filter.
filtered_x = lfilter(taps, 1.0, x)

#------------------------------------------------
# Plot the FIR filter coefficients.
#------------------------------------------------
figure(1)
plot(taps, 'bo-', linewidth=2)
title('Filter Coefficients (%d taps)' % N)
grid(True)

#------------------------------------------------
# Plot the magnitude response of the filter.
#------------------------------------------------
figure(2)
clf()
w, h = freqz(taps, worN=8000)
db = 20*log10(maximum(absolute(h), 1e-5))
plot((w/pi)*nyq_rate, db, linewidth=2)

xlabel('Frequency (Hz)')
ylabel('Gain (db)')
title('Frequency Response')
ylim(-120, 5)
grid(True)

#------------------------------------------------
# Plot the original and filtered signals.
#------------------------------------------------
figure(3)
# Plot the original signal.
subplot(2, 1, 1)
plot(t, x)
title('Original signal')
subplots_adjust(hspace = 0.5)

# Plot the filtered signal, has some phase delay.
subplot(2, 1, 2)
plot(t, filtered_x, 'g', linewidth=4)
title('Filtered signal')

xlabel('t')
grid(True)
show()
